#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
拦截控制API路由
"""

import logging
import json
from datetime import datetime
from bson import ObjectId
from flask import Blueprint, request, jsonify
from db.mongo import get_db
from intercept.services.intercept_service import InterceptService
from intercept.services.iptables_manager import IptablesManager
from intercept.services.proxy_service import ProxyService
from intercept.services.ebpf_capture import EbpfCaptureService

# 尝试导入 libbpf 捕获服务
try:
    from intercept.services.libbpf_capture import LibbpfCaptureService
    LIBBPF_AVAILABLE = True
except ImportError:
    LIBBPF_AVAILABLE = False

# 自定义JSON编码器，处理MongoDB的ObjectId和datetime
class JSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, ObjectId):
            return str(obj)
        if isinstance(obj, datetime):
            return obj.isoformat()
        return json.JSONEncoder.default(self, obj)

logger = logging.getLogger(__name__)
intercept_control_bp = Blueprint('intercept_control', __name__)

@intercept_control_bp.route('/services', methods=['POST'])
def create_service():
    """创建拦截服务"""
    try:
        data = request.json
        if not data:
            return jsonify({"error": "请求体不能为空"}), 400

        # 获取服务信息
        service_name = data.get('service_name')
        service_ip = data.get('service_ip')
        service_port = data.get('service_port')

        if not service_name:
            return jsonify({"error": "服务名称不能为空"}), 400

        # 创建服务
        service = InterceptService.create_intercept_service(
            service_name=service_name,
            service_ip=service_ip,
            service_port=service_port
        )

        if not service:
            return jsonify({"error": "创建服务失败"}), 500

        return jsonify(service), 201

    except Exception as e:
        logger.error(f"创建服务失败: {e}")
        return jsonify({"error": f"创建服务失败: {str(e)}"}), 500

@intercept_control_bp.route('/services', methods=['GET'])
def get_services():
    """获取拦截服务列表"""
    try:
        # 获取服务列表
        services = InterceptService.get_intercept_services()

        return jsonify({
            "services": services,
            "total": len(services)
        }), 200

    except Exception as e:
        logger.error(f"获取服务列表失败: {e}")
        return jsonify({"error": f"获取服务列表失败: {str(e)}"}), 500

@intercept_control_bp.route('/proxy/start', methods=['POST'])
def start_proxy():
    """启动代理服务"""
    try:
        data = request.json or {}
        port = data.get('port', 5555)

        # 启动代理服务
        success = ProxyService.start(port=port)

        if success:
            return jsonify({
                "message": f"代理服务已启动，端口: {port}",
                "port": port,
                "is_running": True
            }), 200
        else:
            return jsonify({
                "message": "代理服务已经在运行",
                "port": ProxyService.get_port(),
                "is_running": True
            }), 200

    except Exception as e:
        logger.error(f"启动代理服务失败: {e}")
        return jsonify({"error": f"启动代理服务失败: {str(e)}"}), 500

@intercept_control_bp.route('/proxy/stop', methods=['POST'])
def stop_proxy():
    """停止代理服务"""
    try:
        # 停止代理服务
        success = ProxyService.stop()

        if success:
            return jsonify({
                "message": "代理服务已停止",
                "is_running": False
            }), 200
        else:
            return jsonify({
                "message": "代理服务未在运行",
                "is_running": False
            }), 200

    except Exception as e:
        logger.error(f"停止代理服务失败: {e}")
        return jsonify({"error": f"停止代理服务失败: {str(e)}"}), 500

@intercept_control_bp.route('/proxy/status', methods=['GET'])
def get_proxy_status():
    """获取代理服务状态"""
    try:
        is_running = ProxyService.is_running()
        port = ProxyService.get_port() if is_running else None

        return jsonify({
            "is_running": is_running,
            "port": port
        }), 200

    except Exception as e:
        logger.error(f"获取代理服务状态失败: {e}")
        return jsonify({"error": f"获取代理服务状态失败: {str(e)}"}), 500

@intercept_control_bp.route('/capture/start', methods=['POST'])
def start_capture():
    """启动捕获服务"""
    try:
        data = request.json or {}
        interface = data.get('interface', 'eth0')
        service_ids = data.get('service_ids', [])

        # 获取服务信息
        services = []
        if service_ids:
            db = get_db()
            for service_id in service_ids:
                service = InterceptService.get_intercept_service_by_id(service_id)
                if service:
                    services.append(service)

        # 启动捕获服务
        # 如果可用，优先使用 libbpf (CO-RE) 方式
        if LIBBPF_AVAILABLE:
            logger.info("使用 libbpf (CO-RE) 方式启动捕获服务")
            try:
                success = LibbpfCaptureService.start_capture(interface=interface, services=services)
                if not success:
                    logger.warning("使用 libbpf 启动捕获服务失败，尝试使用 BCC 模式")
                    success = EbpfCaptureService.start_capture(interface=interface, services=services)
            except Exception as e:
                logger.error(f"使用 libbpf 启动捕获服务异常: {e}")
                logger.warning("尝试使用 BCC 模式")
                success = EbpfCaptureService.start_capture(interface=interface, services=services)
        else:
            # 使用 BCC 或模拟模式
            logger.info("使用 BCC 模式启动捕获服务")
            success = EbpfCaptureService.start_capture(interface=interface, services=services)

        if success:
            return jsonify({
                "message": f"捕获服务已启动，网络接口: {interface}，服务数量: {len(services)}",
                "interface": interface,
                "services": services,
                "is_running": True
            }), 200
        else:
            return jsonify({
                "message": "捕获服务已经在运行",
                "services": EbpfCaptureService.get_services(),
                "is_running": True
            }), 200

    except Exception as e:
        logger.error(f"启动捕获服务失败: {e}")
        return jsonify({"error": f"启动捕获服务失败: {str(e)}"}), 500

@intercept_control_bp.route('/capture/stop', methods=['POST'])
def stop_capture():
    """停止捕获服务"""
    try:
        # 停止捕获服务
        # 如果 libbpf 捕获服务正在运行，则停止它
        if LIBBPF_AVAILABLE and LibbpfCaptureService.is_running():
            logger.info("停止 libbpf 捕获服务")
            success = LibbpfCaptureService.stop_capture()
        else:
            # 停止 BCC 捕获服务
            success = EbpfCaptureService.stop_capture()

        if success:
            return jsonify({
                "message": "捕获服务已停止",
                "is_running": False
            }), 200
        else:
            return jsonify({
                "message": "捕获服务未在运行",
                "is_running": False
            }), 200

    except Exception as e:
        logger.error(f"停止捕获服务失败: {e}")
        return jsonify({"error": f"停止捕获服务失败: {str(e)}"}), 500

@intercept_control_bp.route('/capture/status', methods=['GET'])
def get_capture_status():
    """获取捕获服务状态"""
    try:
        is_running = EbpfCaptureService.is_running()
        services = EbpfCaptureService.get_services() if is_running else []

        return jsonify({
            "is_running": is_running,
            "services": services
        }), 200

    except Exception as e:
        logger.error(f"获取捕获服务状态失败: {e}")
        return jsonify({"error": f"获取捕获服务状态失败: {str(e)}"}), 500

@intercept_control_bp.route('/iptables/setup', methods=['POST'])
def setup_iptables():
    """设置 iptables 规则"""
    try:
        data = request.json
        if not data:
            return jsonify({"error": "请求体不能为空"}), 400

        service_id = data.get('service_id')
        if not service_id:
            return jsonify({"error": "服务ID不能为空"}), 400

        # 获取拦截方向
        direction = data.get('direction', 'output')  # 默认只拦截出站流量
        if direction not in ['input', 'output', 'both']:
            return jsonify({"error": "拦截方向无效，必须是 'input', 'output' 或 'both'"}), 400

        # 获取服务信息
        service = InterceptService.get_intercept_service_by_id(service_id)
        if not service:
            return jsonify({"error": "服务不存在"}), 404

        service_ip = service.get('service_ip')
        service_port = service.get('service_port')
        service_name = service.get('service_name')

        if not service_ip or not service_port:
            return jsonify({"error": "服务IP或端口不能为空"}), 400

        # 设置 iptables 规则
        success = IptablesManager.setup_intercept(
            service_ip=service_ip,
            service_port=service_port,
            comment=service_name,
            direction=direction
        )

        if success:
            # 更新服务状态
            InterceptService.update_intercept_service(
                service_id=service_id,
                is_active=True
            )

            # 根据方向生成消息
            direction_msg = {
                'input': '入站流量',
                'output': '出站流量',
                'both': '入站和出站流量'
            }[direction]

            return jsonify({
                "message": f"已设置拦截规则({direction_msg}): {service_name} ({service_ip}:{service_port})",
                "service": service,
                "direction": direction
            }), 200
        else:
            return jsonify({"error": "设置拦截规则失败"}), 500

    except Exception as e:
        logger.error(f"设置 iptables 规则失败: {e}")
        return jsonify({"error": f"设置 iptables 规则失败: {str(e)}"}), 500

@intercept_control_bp.route('/iptables/remove', methods=['POST'])
def remove_iptables():
    """移除 iptables 规则"""
    try:
        data = request.json
        if not data:
            return jsonify({"error": "请求体不能为空"}), 400

        service_id = data.get('service_id')
        if not service_id:
            return jsonify({"error": "服务ID不能为空"}), 400

        # 获取服务信息
        service = InterceptService.get_intercept_service_by_id(service_id)
        if not service:
            return jsonify({"error": "服务不存在"}), 404

        service_ip = service.get('service_ip')
        service_port = service.get('service_port')
        service_name = service.get('service_name')

        if not service_ip or not service_port:
            return jsonify({"error": "服务IP或端口不能为空"}), 400

        # 移除 iptables 规则
        success = IptablesManager.remove_intercept(
            service_ip=service_ip,
            service_port=service_port,
            comment=service_name
        )

        if success:
            # 更新服务状态
            InterceptService.update_intercept_service(
                service_id=service_id,
                is_active=False
            )

            return jsonify({
                "message": f"已移除拦截规则: {service_name} ({service_ip}:{service_port})",
                "service": service
            }), 200
        else:
            return jsonify({"error": "移除拦截规则失败"}), 500

    except Exception as e:
        logger.error(f"移除 iptables 规则失败: {e}")
        return jsonify({"error": f"移除 iptables 规则失败: {str(e)}"}), 500

@intercept_control_bp.route('/iptables/list', methods=['GET'])
def list_iptables():
    """列出 iptables 规则"""
    try:
        # 列出 iptables 规则
        intercepts = IptablesManager.list_intercepts()

        return jsonify({
            "intercepts": intercepts,
            "count": len(intercepts)
        }), 200

    except Exception as e:
        logger.error(f"列出 iptables 规则失败: {e}")
        return jsonify({"error": f"列出 iptables 规则失败: {str(e)}"}), 500

@intercept_control_bp.route('/iptables/clear', methods=['POST'])
def clear_iptables():
    """清除所有 iptables 规则"""
    try:
        # 清除所有 iptables 规则
        success = IptablesManager.clear_all_intercepts()

        if success:
            # 更新所有服务状态
            db = get_db()
            db.intercept_services.update_many({}, {"$set": {"is_active": False}})

            return jsonify({
                "message": "已清除所有拦截规则"
            }), 200
        else:
            return jsonify({"error": "清除拦截规则失败"}), 500

    except Exception as e:
        logger.error(f"清除 iptables 规则失败: {e}")
        return jsonify({"error": f"清除 iptables 规则失败: {str(e)}"}), 500
